# ==================================================================================
# MaixCAM - 串口遥控的YOLOv8颜色识别与定位系统 (v_protocol_v3_special_display)
# ==================================================================================

#1.导入模块
from maix import camera, display, image, nn, app, uart, pinmap, time # pyright: ignore[reportMissingImports]
import struct

#2.全局配置
UART_DEVICE = "/dev/ttyS1"
UART_BAUD_RATE = 115200
UART_TX_PIN_NAME = "A19"
UART_RX_PIN_NAME = "A18"
UART_FUNCTION = "UART1_TX"
UART_RX_FUNCTION = "UART1_RX"

RX_FRAME_HEADER = 0xAA; RX_FRAME_TAIL = 0xEE; RX_PACKET_LENGTH = 3
TX_FRAME_HEADER = 0xAA; TX_FRAME_TAIL = 0xBB

MODE_IDLE = 0; MODE_1_BLUE_DETECTION = 1; MODE_2_GREEN_DETECTION = 2; MODE_3_RED_DETECTION = 3
MODE_4_BLUE_SPECIAL_DETECTION = 4

COLOR_CLASS_MAPPING = {
    MODE_1_BLUE_DETECTION: [0, 1],
    MODE_2_GREEN_DETECTION: [2, 3],
    MODE_3_RED_DETECTION: [4, 5],
    MODE_4_BLUE_SPECIAL_DETECTION: [0, 1]
}

MODEL_PATH = "/root/gongxun_model/best1.mud"
DETECTION_CONF_THRESHOLD = 0.4
DETECTION_IOU_THRESHOLD = 0.45

DRAW_COLORS = {
    "blue": image.Color(0, 0, 255), "green": image.Color(0, 255, 0),
    "red": image.Color(255, 0, 0), "text": image.Color(255, 255, 255),
    "center": image.Color(255, 255, 0)
}

#3.初始化
current_mode = MODE_IDLE; serial_port = None; cam = None; disp = None
detector = None; initialization_successful = False; rx_buffer = bytearray()
screen_center_x = 0; screen_center_y = 0

try:
    pinmap.set_pin_function(UART_TX_PIN_NAME, UART_FUNCTION)
    pinmap.set_pin_function(UART_RX_PIN_NAME, UART_RX_FUNCTION)
    serial_port = uart.UART(UART_DEVICE, UART_BAUD_RATE)
    print(f"UART 设备 {UART_DEVICE} @ {UART_BAUD_RATE} 初始化成功")

    detector = nn.YOLO11(model=MODEL_PATH, dual_buff=True)
    print(f"模型 {MODEL_PATH} 加载成功. 标签: {detector.labels}")

    cam_width = detector.input_width(); cam_height = detector.input_height()
    cam = camera.Camera(cam_width, cam_height, detector.input_format())
    
    screen_center_x = cam_width // 2; screen_center_y = cam_height // 2
    print(f"屏幕尺寸: {cam_width}x{cam_height}, 中心点: ({screen_center_x}, {screen_center_y})")

    disp = display.Display()
    print("显示屏初始化成功")

    try: cam.skip_frames(30)
    except Exception: time.sleep_ms(1000)
        
    initialization_successful = True
    print("系统初始化成功")

except Exception as e:
    print(f"初始化失败: {e}"); app.set_exit_flag(True)

#4.辅助函数
def process_serial_data():
    global current_mode, rx_buffer
    if not serial_port: return
    read_data = serial_port.read(128)
    if read_data: rx_buffer.extend(read_data)
    while len(rx_buffer) >= RX_PACKET_LENGTH:
        if rx_buffer[0] == RX_FRAME_HEADER and rx_buffer[RX_PACKET_LENGTH - 1] == RX_FRAME_TAIL:
            command = rx_buffer[1]
            if command == 0x01: current_mode = MODE_1_BLUE_DETECTION
            elif command == 0x02: current_mode = MODE_2_GREEN_DETECTION
            elif command == 0x03: current_mode = MODE_3_RED_DETECTION
            elif command == 0x04: current_mode = MODE_4_BLUE_SPECIAL_DETECTION
            elif command == 0x00: current_mode = MODE_IDLE
            else: print(f"接收到未知指令: {hex(command)}")
            rx_buffer = rx_buffer[RX_PACKET_LENGTH:]
        else:
            rx_buffer.pop(0)

def send_data(serial, class_id, val1, val2):
    if not serial: return
    try:
        packet = struct.pack('<BBhhB', TX_FRAME_HEADER, int(class_id), int(val1), int(val2), TX_FRAME_TAIL)
        serial.write(packet)
    except Exception as e:
        print(f"串口发送错误: {e}")

#主循环
if initialization_successful:
    last_mode = -1
    try:
        while not app.need_exit():
            process_serial_data()
            img = cam.read()
            if not img: continue

            target_obj = None
            dx, dy = "N/A", "N/A"

            if current_mode != MODE_IDLE:
                objs = detector.detect(img, conf_th=DETECTION_CONF_THRESHOLD, iou_th=DETECTION_IOU_THRESHOLD)
                
                target_classes = COLOR_CLASS_MAPPING.get(current_mode, [])
                valid_objs = [obj for obj in objs if obj.class_id in target_classes]
                
                if valid_objs:
                    target_obj = max(valid_objs, key=lambda o: o.w * o.h)

            if target_obj:
                if current_mode in [MODE_1_BLUE_DETECTION, MODE_2_GREEN_DETECTION, MODE_3_RED_DETECTION]:
                    obj_center_x = target_obj.x + target_obj.w // 2
                    obj_center_y = target_obj.y + target_obj.h // 2
                    dx = obj_center_x - screen_center_x
                    dy = obj_center_y - screen_center_y
                    send_data(serial_port, current_mode, obj_center_x, obj_center_y)
                
                elif current_mode == MODE_4_BLUE_SPECIAL_DETECTION:
                    val1 = target_obj.x 
                    val2 = target_obj.y
                    send_data(serial_port, current_mode, val1, val2)
                    
                    #绘制特殊计算的点
                    img.draw_cross(val1, val2, size=15, color=image.COLOR_YELLOW, thickness=2)
                    
                    obj_center_x = target_obj.x + target_obj.w // 2
                    obj_center_y = target_obj.y + target_obj.h // 2
                    dx = obj_center_x - screen_center_x
                    dy = obj_center_y - screen_center_y

                color_name = detector.labels[target_obj.class_id].rstrip('0123456789')
                box_color = DRAW_COLORS.get(color_name, DRAW_COLORS["text"])
                img.draw_rect(target_obj.x, target_obj.y, target_obj.w, target_obj.h, color=box_color, thickness=2)
                msg = f'{detector.labels[target_obj.class_id]}: {target_obj.score:.2f}'
                img.draw_string(target_obj.x, target_obj.y - 15, msg, color=box_color, scale=1.2)
            
            mode_str = "IDLE"
            if current_mode != last_mode:
                if current_mode in COLOR_CLASS_MAPPING:
                    color_name = detector.labels[COLOR_CLASS_MAPPING[current_mode][0]].rstrip('0123456789')
                    if current_mode == MODE_4_BLUE_SPECIAL_DETECTION:
                        mode_str = "DETECT (Blue Special)"
                    else:
                        mode_str = f"DETECT ({color_name})"
                    print(f"切换到模式: {mode_str}")
                else:
                    mode_str = "IDLE"
                    print("切换到模式: IDLE")
                last_mode = current_mode
            else:
                if current_mode in COLOR_CLASS_MAPPING:
                    color_name = detector.labels[COLOR_CLASS_MAPPING[current_mode][0]].rstrip('0123456789')
                    if current_mode == MODE_4_BLUE_SPECIAL_DETECTION:
                        mode_str = "DETECT (Blue Special)"
                    else:
                        mode_str = f"DETECT ({color_name})"
                else:
                    mode_str = "IDLE"

            img.draw_string(5, 5, f"Mode: {mode_str}", color=DRAW_COLORS["text"], scale=1.5)
            img.draw_string(img.width() - 100, 5, f"FPS:{time.fps():.1f}", color=DRAW_COLORS["text"], scale=1.5)
            img.draw_string(5, 30, f"dx: {dx}", color=DRAW_COLORS["text"], scale=1.5)
            img.draw_string(5, 55, f"dy: {dy}", color=DRAW_COLORS["text"], scale=1.5)
            img.draw_cross(screen_center_x, screen_center_y, size=10, color=image.COLOR_RED, thickness=2)
            
            disp.show(img)

    except Exception as e:
        print(f"主循环出错: {e}")
    finally:
        print("程序退出。")
else:
    print("错误：初始化失败，无法运行主循环。")